package com.lee.windowFunction;

import com.lee.IP;
import com.lee.TCP;
import com.lee.UDP;
import com.lee.networkData;
import org.apache.flink.streaming.api.functions.windowing.AllWindowFunction;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;

/**
 * Created by liruopeng on 2018/1/18.
 */
public class StatisticalCountWindowFunction implements AllWindowFunction<networkData, networkData, TimeWindow> {
    private Logger logger = LoggerFactory.getLogger(StatisticalCountWindowFunction.class);

    private HashMap<String, Integer> countDstMap;
    private HashMap<String, Integer> countServiceMap;


    public StatisticalCountWindowFunction(HashMap<String, Integer> countDstMap, HashMap<String, Integer> countServiceMap) {
        this.countDstMap = countDstMap;
        this.countServiceMap = countServiceMap;
    }

    @Override
    public void apply(TimeWindow window, Iterable<networkData> values, Collector<networkData> out) throws Exception {
        int totalDstNumber = 0, totalServiceNumber = 0;
        for (String dst : countDstMap.keySet()) {
            totalDstNumber += countDstMap.get(dst);
        }
        for (String service : countServiceMap.keySet()) {
            totalDstNumber += countServiceMap.get(service);
        }
        for (networkData value : values) {
            if (!(value.getInternetLayer() instanceof IP)) {
                // 如果不是IP协议 不用计算，直接输入值
                value.setCountWindowSameDstCount(0);
                value.setCountWindowSameDstPercent(0);
                value.setCountWindowDiffDstCount(totalDstNumber);
                value.setCountWindowDiffDstPercent(1);
                value.setCountWindowSameServiceCount(0);
                out.collect(value);
                continue;
            }

            IP ip = (IP) value.getInternetLayer();
            if (countDstMap.containsKey(ip.getDstIP().toString())) {
                // 一般一定在相同地址中含有
                int count = countDstMap.get(ip.getDstIP().toString());
                value.setCountWindowSameDstCount(count);
                value.setCountWindowSameDstPercent(((float) count) / ((float) totalDstNumber));
                value.setCountWindowDiffDstCount((totalDstNumber - count));
                value.setCountWindowDiffDstPercent(((float) (totalDstNumber - count)) / ((float) totalDstNumber));

            } else {
                // 如果发生了这种情况，说明计算出错了
                logger.error("目的地址map中不含有当前IP地址，请检查代码错误！！");
                value.setCountWindowSameDstCount(0);
                value.setCountWindowSameDstPercent(0);
                value.setCountWindowDiffDstCount(totalDstNumber);
                value.setCountWindowDiffDstPercent(1);
                value.setCountWindowSameServiceCount(0);
            }
            if (value.getTransportLayer() == null) {
                // 不属于TCP和UDP协议，直接赋予值
                value.setCountWindowSameServiceCount(0);
                out.collect(value);
                continue;
            }
            if (value.getTransportLayer() instanceof TCP) {
                // 传输层是TCP协议
                TCP tcp = (TCP) value.getTransportLayer();
                if (tcp.getApplicationProtocol().toString().compareTo("unknown") == 0) {
                    setCountWindowSameServiceNumberTCP(value, tcp);
                } else {
                    setCountWindowSameServiceNumberTCP(value, tcp);
                }
            } else {
                // 传输层是UDP协议
                UDP udp = (UDP) value.getTransportLayer();
                if (udp.getApplicationProtocol().toString().compareTo("unknown") == 0) {
                    setCountWindowSameServiceNumberUDP(value, udp);
                } else {
                    setCountWindowSameServiceNumberUDP(value, udp);
                }
            }
            out.collect(value);
        }

        countDstMap.clear();
        countServiceMap.clear();
    }

    public void setCountWindowSameServiceNumberTCP(networkData value, TCP tcp) {
        if (countServiceMap.containsKey(tcp.getDstPort().toString())) {
            value.setTimeWindowSameServiceCount(countServiceMap.get(tcp.getDstPort().toString()));
        } else {
            value.setTimeWindowSameServiceCount(0);
        }
    }

    public void setCountWindowSameServiceNumberUDP(networkData value, UDP udp) {
        if (countServiceMap.containsKey(udp.getDstPort().toString())) {
            value.setTimeWindowSameServiceCount(countServiceMap.get(udp.getDstPort().toString()));
        } else {
            value.setTimeWindowSameServiceCount(0);
        }
    }
}